import torch.nn as nn

class SeparableConv1D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding='valid'):
        super(SeparableConv1D, self).__init__()
        # Depthwise convolution
        self.depthwise = nn.Conv1d(in_channels, in_channels, kernel_size, stride, padding=0 if padding == 'valid' else 'same', groups=in_channels)
        # Pointwise convolution
        self.pointwise = nn.Conv1d(in_channels, out_channels, 1, 1)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

class SepConv1DNet(nn.Module):
    def __init__(self):
        super().__init__()
        # ZeroPadding1D
        self.zero_pad = nn.ZeroPad1d(padding=4)
        # SeparableConv1D
        # self.sep_conv1d = nn.Conv1d(in_channels=6, out_channels=4, kernel_size=16, stride=8, groups=6)
        self.sep_conv1d = SeparableConv1D(in_channels=6,out_channels=4,kernel_size=16,stride=8)
        self.activation = nn.Tanh()
        # Flatten
        self.dense = nn.Linear(in_features=38 * 4, out_features=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.zero_pad(x)
        x = self.sep_conv1d(x)
        x = self.activation(x)
        x = x.view(x.size(0), -1)
        x = self.dense(x)
        x = self.sigmoid(x)
        return x

'''
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
         ZeroPad1d-1               [-1, 6, 315]               0
            Conv1d-2                [-1, 6, 38]             102
            Conv1d-3                [-1, 4, 38]              28
   SeparableConv1D-4                [-1, 4, 38]               0
              Tanh-5                [-1, 4, 38]               0
            Linear-6                    [-1, 1]             153
           Sigmoid-7                    [-1, 1]               0
================================================================
Total params: 283
Trainable params: 283
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.01
Forward/backward pass size (MB): 0.02
Params size (MB): 0.00
Estimated Total Size (MB): 0.03
----------------------------------------------------------------
None
'''